import os
import xml.etree.ElementTree as ET
from xml.etree.ElementTree import Element
#/home/aiserver/keras-yolo3-master/新建文件夹/ftp/xml
#需要处理的文件夹目录的绝对路径
IMG_ROOT_DIR="/home/aiserver/keras-yolo3-master/VOC2007/JPEGImages/"
#修改的标签值
#LABEL="trashbin"

#是否要修改标签
IS_SET_LABEL=True

#替换标签字典
EX_LABEL_DIC={

"hat":"hat_",


}

#处理目录
PROCESS_IMAGE_DIR= r"./VOCdevkit/VOC/Annotations/"  #注意最后的\\

#格式控制
FORMAT_SET=(".xml")

#标签过滤
LABEL_FILTER=("face")


#是否进行标签类型过滤
IS_FILET_LABEL=True

def write_xml(tree, out_path):
    '''''将xml文件写出
       tree: xml树
       out_path: 写出路径'''
    tree.write(out_path, encoding="utf-8",xml_declaration=False)
    print("写入成功》》》")

def process_xml(file,filename_):
    tree = ET.parse(file)
    #print(type(tree))
    root_ele = tree.getroot()  # 使用getroot()获取根节点，得到的是一个Element对象
    #print(type(root_ele))
    #object_root=root.find("object")
    #path=root.find("filename")
    #JPEGImages
    root_ele.find("filename").text=filename_
    #print(filename)
    filename_all=filename_+".xml"
    #print(filename_all)
    new_path=IMG_ROOT_DIR+filename_+".jpg"
    #print(new_path)
    path_ele=root_ele.find("path")
    #print(">>>>")
    #print(path_ele)


    #添加《path》节点
    if path_ele==None:
        path_ele = Element("path")
        path_ele.text = new_path#=IMG_ROOT_DIR+filename  # 二级目录的值
        root_ele.append(path_ele)  # 将二级目录加到一级目录里
    path_ele.text = new_path
    #path_ele.text=new_path
    #print("<>>>>>>")
    #print(new_path)
    #print(path_ele.text)
    folder_ele = root_ele.find("folder")
    folder_ele.text = "JPEGImages"
    #print(folder_ele.text)

    object_all_ele=root_ele.findall("object")
    #标准化《filename》
    #path_ele = root_ele.find("filename")

    #print(">>>>")
    #print(path_ele)

    #修改标签名称
    if IS_SET_LABEL:
        for object_ in object_all_ele:
            #print("???")
            print("修改标签成功")
            object_name_ele=object_.find("name")
            current_label=object_name_ele.text
            if current_label in list(EX_LABEL_DIC.keys()):
                object_name_ele.text=EX_LABEL_DIC[current_label]
            #print(object_.find("name").text)
            #print(path_ele.text)

    #进行标签过滤
    if IS_FILET_LABEL:
        for object_ in object_all_ele:
            #print("???")
            object_name_ele=object_.find("name")
            object_name=object_name_ele.text
            #print(object_name)
            if object_name not in LABEL_FILTER:
                root_ele.remove(object_)
                print("删除节点成功")
                #print(object_)
            #print(object_.find("name").text)
            #print(path_ele.text)
    #tree.write(new_path,encoding="utf-8")
    #print("WRITING>>>>>")

    storage_dir=PROCESS_IMAGE_DIR+filename_all
    #print(">>>")
    #print(storage_dir)
    #print(PROCESS_IMAGE_DIR)
    #print(filename_all)
    write_xml(tree,storage_dir)
    print(storage_dir)

def getfiles(PROCESS_IMAGE_DIR):
    filenames=os.listdir(PROCESS_IMAGE_DIR)
    #i=157

    for file in filenames:
        #print(file)
        filename_=os.path.splitext(file)[0]
        format = os.path.splitext(file)[1]
        if(format in FORMAT_SET):
            #print(format)
            img_dir = os.path.join(PROCESS_IMAGE_DIR, file)
            process_xml(img_dir,filename_)
if __name__ == '__main__':
    getfiles(PROCESS_IMAGE_DIR)
